/*
 * Copyright (c) 2020-2099 the original author or authors. All rights reserve.
 */
package org.tianyun.cloud.utils;

import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import javax.crypto.Cipher;
import javax.crypto.KeyGenerator;
import javax.crypto.SecretKey;
import javax.crypto.spec.IvParameterSpec;
import javax.crypto.spec.SecretKeySpec;
import java.security.Key;
import java.security.KeyFactory;
import java.security.KeyPair;
import java.security.KeyPairGenerator;
import java.security.NoSuchAlgorithmException;
import java.security.PrivateKey;
import java.security.PublicKey;
import java.security.SecureRandom;
import java.security.Security;
import java.security.Signature;
import java.security.spec.InvalidKeySpecException;
import java.security.spec.X509EncodedKeySpec;

/**
 * 加解密算法
 *
 * @auther ebert_chan
 */
public class CryptBuilder {

    private static final Logger LOGGER = LoggerFactory.getLogger(CryptBuilder.class);

    static {
        Security.addProvider(new org.bouncycastle.jce.provider.BouncyCastleProvider());
    }

    private CryptBuilder() {
    }

    /**
     * rsa加解密算法，keySize：1024
     *
     * @return
     * @auther ebert_chan
     */
    public static RsaAlgorithm rsa() {
        return new RsaAlgorithm();
    }

    /**
     * rsa加解密算法，keySize：2048
     *
     * @return
     * @auther ebert_chan
     */
    public static RsaAlgorithm rsa2048() {
        RsaAlgorithm rsaAlgorithm = new RsaAlgorithm();
        rsaAlgorithm.keySize(2048);
        return rsaAlgorithm;
    }

    /**
     * aes加解密算法，keySize：128
     *
     * @return
     * @auther ebert_chan
     */
    public static AesAlgorithm aes() {
        return new AesAlgorithm();
    }

    /**
     * aes加解密算法，keySize：192
     *
     * @return
     * @auther ebert_chan
     */
    public static AesAlgorithm aes192() {
        AesAlgorithm aesAlgorithm = new AesAlgorithm();
        aesAlgorithm.init(AesAlgorithm.AES_ALGORITHM, 192);
        return aesAlgorithm;
    }

    /**
     * aes加解密算法，keySize：256
     *
     * @return
     * @auther ebert_chan
     */
    public static AesAlgorithm aes256() {
        AesAlgorithm aesAlgorithm = new AesAlgorithm();
        aesAlgorithm.init(AesAlgorithm.AES_ALGORITHM, 256);
        return aesAlgorithm;
    }

    /**
     * rsa加解密算法，默认密匙长度1024
     *
     * @auther ebert_chan
     */
    public static class RsaAlgorithm {

        private static final String RSA_ALGORITHM = "RSA";

        private static final String SIGNATURE_ALGORITHM = "MD5withRSA";

        private String rsaAlgorithm = RSA_ALGORITHM;

        private String signatureAlgorithm = SIGNATURE_ALGORITHM;

        private int keySize = 1024;

        private PublicKey publicKey;

        private PrivateKey privateKey;

        private RsaAlgorithm() {
        }

        private RsaAlgorithm keySize(int keySize) {
            this.keySize = keySize;
            return this;
        }

        public RsaAlgorithm generateKeyPair() {
            // 使用RSA算法获得密钥对生成器对象 keyPairGenerator
            KeyPairGenerator keyPairGenerator;
            try {
                keyPairGenerator = KeyPairGenerator.getInstance(rsaAlgorithm);
                // 设置密钥长度
                keyPairGenerator.initialize(keySize);
                // 生成密钥对
                KeyPair keyPair = keyPairGenerator.generateKeyPair();
                // 获取公钥
                publicKey = keyPair.getPublic();
                // 获取私钥
                privateKey = keyPair.getPrivate();
            } catch (NoSuchAlgorithmException e) {
                LOGGER.warn("rsa no such algorithm. {}", e.getMessage());
            }
            return this;
        }

        public String getPublicKey() {
            if (publicKey == null) {
                return null;
            }
            byte[] encoded = publicKey.getEncoded();
            return CryptHelper.base64EncryptInline(encoded);
        }

        public RsaAlgorithm publicKey(byte[] encodedKey) {
            if (encodedKey == null) {
                return this;
            }

            try {
                // 构造X509EncodedKeySpec对象
                X509EncodedKeySpec keySpec = new X509EncodedKeySpec(encodedKey);
                // ALGORITHM 指定的加密算法
                KeyFactory keyFactory;
                keyFactory = KeyFactory.getInstance(rsaAlgorithm);
                // 取公匙对象
                publicKey = keyFactory.generatePublic(keySpec);
            } catch (NoSuchAlgorithmException e) {
                LOGGER.warn("rsa no such algorithm. {}", e.getMessage());
            } catch (InvalidKeySpecException e) {
                LOGGER.warn("rsa invalid key. {}", e.getMessage());
            }
            return this;
        }

        public RsaAlgorithm publicKey(String publicKeyStr) {
            byte[] encodedKey = CryptHelper.base64DecryptToBuffer(publicKeyStr);
            this.publicKey(encodedKey);
            return this;
        }

        public String getPrivateKey() {
            if (privateKey == null) {
                return null;
            }
            byte[] encoded = privateKey.getEncoded();
            return CryptHelper.base64EncryptInline(encoded);
        }

        public RsaAlgorithm privateKey(byte[] encodedKey) {
            if (encodedKey == null) {
                return this;
            }

            try {
                // 构造X509EncodedKeySpec对象
                X509EncodedKeySpec keySpec = new X509EncodedKeySpec(encodedKey);
                // ALGORITHM 指定的加密算法
                KeyFactory keyFactory;
                keyFactory = KeyFactory.getInstance(rsaAlgorithm);
                // 取私匙对象
                privateKey = keyFactory.generatePrivate(keySpec);
            } catch (NoSuchAlgorithmException e) {
                LOGGER.warn("rsa no such algorithm. {}", e.getMessage());
            } catch (InvalidKeySpecException e) {
                LOGGER.warn("rsa invalid key. {}", e.getMessage());
            }
            return this;
        }

        public RsaAlgorithm privateKey(String privateKeyStr) {
            byte[] encodedKey = CryptHelper.base64DecryptToBuffer(privateKeyStr);
            this.privateKey(encodedKey);
            return this;
        }

        public String sign(byte[] dataBuffer) {
            if (privateKey == null || dataBuffer == null) {
                return null;
            }
            Signature signet;
            try {
                signet = Signature.getInstance(signatureAlgorithm);
                signet.initSign(privateKey);
                signet.update(dataBuffer);
                byte[] signature = signet.sign();
                return CryptHelper.base64EncryptInline(signature);
            } catch (Throwable e) {
                LOGGER.warn("rsa sign failed", e.getMessage());
            }
            return null;
        }

        public String sign(String data) {
            byte[] dataBuffer = CryptHelper.stringToBytes(data);
            return this.sign(dataBuffer);
        }

        public boolean verify(byte[] dataBuffer, byte[] signature) {
            if (publicKey == null || signature == null || dataBuffer == null) {
                return false;
            }
            Signature signet;
            try {
                signet = Signature.getInstance(signatureAlgorithm);
                signet.initVerify(publicKey);
                signet.update(dataBuffer);
                return signet.verify(signature);
            } catch (Throwable e) {
                LOGGER.warn("rsa verify failed", e.getMessage());
            }
            return false;
        }

        public boolean verify(String data, String sign) {
            byte[] dataBuffer = CryptHelper.stringToBytes(data);
            byte[] signature = CryptHelper.base64DecryptToBuffer(sign);
            return this.verify(dataBuffer, signature);
        }

        private String encrypt(byte[] dataBuffer, Key key) {
            if (dataBuffer == null || key == null) {
                return null;
            }
            try {
                // 获取一个加密算法为RSA的加解密器对象cipher。
                Cipher cipher = Cipher.getInstance(rsaAlgorithm);
                // 设置为加密模式,并将密匙给cipher。
                cipher.init(Cipher.ENCRYPT_MODE, key);
                byte[] result = cipher.doFinal(dataBuffer);
                return CryptHelper.base64EncryptInline(result);
            } catch (Throwable e) {
                LOGGER.warn("aes encrypt failed. {}", e.getMessage());
            }
            return null;
        }

        private String encrypt(String data, Key key) {
            byte[] dataBuffer = CryptHelper.stringToBytes(data);
            return this.encrypt(dataBuffer, key);
        }

        public String encryptByPublicKey(byte[] dataBuffer) {
            return this.encrypt(dataBuffer, publicKey);
        }

        public String encryptByPublicKey(String data) {
            return this.encrypt(data, publicKey);
        }

        public String encryptByPrivateKey(byte[] dataBuffer) {
            return this.encrypt(dataBuffer, privateKey);
        }

        public String encryptByPrivateKey(String data) {
            return this.encrypt(data, privateKey);
        }

        private byte[] decrypt(byte[] dataBuffer, Key key) {
            if (dataBuffer == null || key == null) {
                return null;
            }
            try {
                // 获取一个加密算法为RSA的加解密器对象cipher。
                Cipher cipher = Cipher.getInstance(rsaAlgorithm);
                // 设置为解密模式,并将密匙给cipher。
                cipher.init(Cipher.DECRYPT_MODE, key);
                byte[] result = cipher.doFinal(dataBuffer);
                return result;
            } catch (Throwable e) {
                LOGGER.warn("aes decrypt failed. {}", e.getMessage());
            }
            return null;
        }

        private String decrypt(String data, Key key) {
            byte[] dataBuffer = CryptHelper.base64DecryptToBuffer(data);
            byte[] result = this.decrypt(dataBuffer, key);
            return CryptHelper.bytesToString(result);
        }

        public byte[] decryptByPublicKey(byte[] dataBuffer) {
            return this.decrypt(dataBuffer, publicKey);
        }

        public String decryptByPublicKey(String data) {
            return this.decrypt(data, publicKey);
        }

        public byte[] decryptByPrivateKey(byte[] dataBuffer) {
            return this.decrypt(dataBuffer, privateKey);
        }

        public String decryptByPrivateKey(String data) {
            return this.decrypt(data, privateKey);
        }

    }

    /**
     * aes加解密，默认对称加密使用的算法为 AES-128-CBC，数据采用PKCS#7填充。
     *
     * @auther ebert_chan
     */
    public static class AesAlgorithm {

        private static final String AES_ALGORITHM = "AES/CBC/PKCS7Padding";

        private static final String AES_ALGORITHM_SHORT = "AES";

        private static final String KEYGEN_ALGORITHM = "SHA1PRNG";

        private String algorithm;

        private int keySize;

        private AesAlgorithm() {
            this.init(AES_ALGORITHM, 128);
        }

        /**
         * 初始化算法
         *
         * @param algorithm 算法，如：AES/CBC/PKCS7Padding
         * @param keySize   初始key大小，如：128
         * @return
         * @auther ebert_chan
         */
        public AesAlgorithm init(String algorithm, int keySize) {
            this.algorithm = algorithm;
            this.keySize = keySize;
            return this;
        }

        /**
         * 加密
         *
         * @param text 明文
         * @param key  秘钥，keySize/8 个字节
         * @param iv   向量，16字节
         * @return
         * @auther ebert_chan
         */
        public String encrypt(String text, String key, String iv) {
            byte[] data = CryptHelper.stringToBytes(text);
            if (data == null) {
                return null;
            }
            byte[] result = encrypt(data, CryptHelper.base64DecryptToBuffer(key), CryptHelper.base64DecryptToBuffer(iv));
            return CryptHelper.base64EncryptInline(result);
        }

        /**
         * 加密
         *
         * @param data 明文
         * @param key  秘钥，keySize/8 个字节
         * @param iv   向量，16字节
         * @return
         * @auther ebert_chan
         */
        public byte[] encrypt(byte[] data, byte[] key, byte[] iv) {
            try {
                SecretKeySpec secretKeySpec = null;
                if (key.length == keySize / 8) {
                    secretKeySpec = new SecretKeySpec(key, AES_ALGORITHM_SHORT);
                } else {
                    KeyGenerator kgen = KeyGenerator.getInstance(AES_ALGORITHM_SHORT);
                    SecureRandom secureRandom = SecureRandom.getInstance(KEYGEN_ALGORITHM);
                    secureRandom.setSeed(key);
                    kgen.init(keySize, secureRandom);
                    SecretKey secretKey = kgen.generateKey();
                    byte[] enCodeFormat = secretKey.getEncoded();
                    secretKeySpec = new SecretKeySpec(enCodeFormat, AES_ALGORITHM_SHORT);
                }

                // 算法/模式/补码方式
                Cipher cipher = Cipher.getInstance(algorithm);
                if (iv != null) {
                    IvParameterSpec ivParameterSpec = new IvParameterSpec(iv);
                    cipher.init(Cipher.ENCRYPT_MODE, secretKeySpec, ivParameterSpec);
                } else {
                    cipher.init(Cipher.ENCRYPT_MODE, secretKeySpec);
                }

                return cipher.doFinal(data);

            } catch (Exception e) {
                LOGGER.warn("aes encrypt failed. {}", e.getMessage());
            }
            return null;
        }

        /**
         * 解密
         *
         * @param text 密文
         * @param key  秘钥，keySize/8 个字节
         * @param iv   向量，16字节
         * @return
         * @auther ebert_chan
         */
        public String decrypt(String text, String key, String iv) {
            byte[] buffer = decrypt(CryptHelper.base64DecryptToBuffer(text), CryptHelper.base64DecryptToBuffer(key), CryptHelper.base64DecryptToBuffer(iv));
            return CryptHelper.bytesToString(buffer);
        }

        /**
         * 解密
         *
         * @param data 密文
         * @param key  秘钥，keySize/8 个字节
         * @param iv   向量，16字节
         * @return
         * @auther ebert_chan
         */
        public byte[] decrypt(byte[] data, byte[] key, byte[] iv) {
            if (data == null || key == null) {
                return null;
            }
            try {
                SecretKeySpec secretKeySpec = null;
                if (key.length == keySize / 8) {
                    secretKeySpec = new SecretKeySpec(key, AES_ALGORITHM_SHORT);
                } else {
                    KeyGenerator kgen = KeyGenerator.getInstance(AES_ALGORITHM_SHORT);
                    SecureRandom secureRandom = SecureRandom.getInstance(KEYGEN_ALGORITHM);
                    secureRandom.setSeed(key);
                    kgen.init(keySize, secureRandom);
                    SecretKey secretKey = kgen.generateKey();
                    byte[] enCodeFormat = secretKey.getEncoded();
                    secretKeySpec = new SecretKeySpec(enCodeFormat, AES_ALGORITHM_SHORT);
                }

                // 算法/模式/补码方式
                Cipher cipher = Cipher.getInstance(algorithm);
                if (iv != null) {
                    IvParameterSpec ivParameterSpec = new IvParameterSpec(iv);
                    cipher.init(Cipher.DECRYPT_MODE, secretKeySpec, ivParameterSpec);
                } else {
                    cipher.init(Cipher.DECRYPT_MODE, secretKeySpec);
                }

                byte[] result = cipher.doFinal(data);

                return result;

            } catch (Exception e) {
                LOGGER.warn("aes decrypt failed. {}", e.getMessage());
            }
            return null;
        }

    }

}
